"""
回测相关的类和函数，
包括编写策略时的模板， 回测函数， 搜参函数， 滚动回测
"""
import copy
import os.path
import random
import datetime
from warnings import filterwarnings
from shutil import rmtree
from multiprocessing import Pool, Queue
from os import cpu_count
from itertools import product

import pandas as pd
import numpy as np

from dytz.data import *
from dytz.bar import *
from dytz.context import Context
from dytz.ga import *
from dytz.params import *

__all__ = ['DNGNode', 'backtest', 'rolling_back', 'rolling_back2', 'select_params', 'select_params2', 'summarization_data']

filterwarnings('ignore')

class DNGNode:
    """有向图模板类， 编写策略时需要继承该基类"""
    __params__ = {

    }

    def __init__(self, **params):
        """
        构造器
        :param params: 参数
        """
        self.params = params
        self.init_dir = dir(self)

    def init_node(self) -> None:
        """
        初始化,主要定义需要使用的时间序列
        :return: None
        """
        pass

    def on_bar(self, ctx: Context, bar: Bar) -> None:
        """
        对每根Bar线进行计算
        :param ctx: 上下文收集器
        :param bar: Bar线
        :return: None
        """
        pass

    def collect(self) -> None:
        """
        收集init_node中创建的时间序列
        :return:
        """
        new_dir = dir(self)
        self.collect_dir = []
        for i in new_dir:
            if i not in self.init_dir:
                if i not in ['init_dir', 'collect_dir', '__slotnames__']:
                    self.collect_dir.append(eval(f'self.{i}'))

    def step(self) -> None:
        """
        对init_node中创建的时间序列执行step()方法
        :return: None
        """
        for i in self.collect_dir:
            i.step()

class DNGIndex(DNGNode):
    def on_bar(self, ctx: Context, bar: Bar):
        ctx.open_long(price=bar.AVGPRICE)

model_ind = [DNGIndex()]


def backtest(models: list[DNGNode, ...], code: str, st: str, et: str,
             plot: bool = False, file_name: str = None, res_type: str = 'score') -> float or dict:
    """
    回测系统
    :param models: 以DNGNode列表组成的策略
    :param code:标的代码
    :param st: 开始时间
    :param et: 结束时间
    :param plot: 是否要画出净值图
    :param file_name: 保存的文件名
    :param res_type: 返回的结果类型，score是对策略打分，目前是shape和clamar的平均值，也可是avg,std,shape,maxDD和calmar的字典
    :return:float or dict， 由res_type指定
    """
    da = iter(ChoicesData(code, st, et))
    ctx = Context()
    for m in models:
        m.init_node()
    for bar in da:
        ctx.on_bar(bar)
        for m in models:
            m.on_bar(ctx, bar)
        ctx.on_sta(bar)
    if res_type == 'score':
        res = ctx.get_score(plot=plot, file_name=file_name)
    elif res_type == 'res':
        res = ctx.get_res(plot=plot, file_name=file_name)
        # if pd.isna(res['avg']):
        #     print(models[0].params)
    else:
        raise
    return res


def __generate_individual(p: dict, models: list[DNGNode, ...],
                          code: str, st: str, et: str) -> OneBack:
    """
    有策略生成个体， 为多进程的原子操作之一
    :param p: 参数
    :param models: 模块
    :param code: 回测标的
    :param st: 开始时间
    :param et: 结束时间
    :return: 回测个体
    """
    one = OneBack(p)
    models[0].params = one.params
    one.res_val = backtest(models, code, st, et)
    # backtest(models, code, st, et, res_type='res')
    return one


def select_params(models: list[DNGNode, ...], code: str, st: str, et: str, n: int = 400) -> None:
    """
    搜参系统， 采用遗传算法
    :param models: 以DNGNode列表组成的策略
    :param code: 回测标的代码
    :param st: 开始时间
    :param et: 结束时间
    :param n:运行轮次
    :return: None
    """
    #  初始种群
    zh_all = int(n * 0.25)
    zhongquan = int(zh_all * 0.5)
    pop = []
    p = models[0].__params__
    pool = Pool(cpu_count() - 1)
    for _ in range(zh_all):  # 构建个体
        pop.append(pool.apply_async(func=__generate_individual, args=(p, models, code, st, et)))
    pool.close()
    pool.join()
    pop = list(map(lambda x: x.get(), pop))
    pool.terminate()
    pop.sort(key=lambda x: x.res_val, reverse=True)
    pop = pop[: zhongquan]
    print(pop[0].res_val, pop[-1].res_val)
    ind = list(range(len(pop)))
    for it in range(n):
        print(it)
        se_i = list(selection(ind, 2))
        se = [pop[se_i[0]], pop[se_i[1]]]
        se_all = copy.deepcopy(se)
        if random.uniform(0, 1) < 0.8:
            p1, p2 = se[0] + se[1]
            models[0].params = p1.params
            p1.res_val = backtest(models, code, st, et)
            se_all += [p1]
            models[0].params = p2.params
            p2.res_val = backtest(models, code, st, et)
            se_all += [p2]
        for i in se:
            if i.aber(0.15):
                models[0].params = i.params
                i.res_val = backtest(models, code, st, et)
                se_all += [i]
        se_all.sort(key=lambda x: x.res_val, reverse=True)
        pop[se_i[0]], pop[se_i[1]] = se_all[0], se_all[1]
    pop.sort(key=lambda x: x.res_val, reverse=True)
    best_all = pop[:20]
    par_all = []
    res_all = []
    for i in range(len(best_all)):
        p_one = best_all[i].params
        par_all.append(p_one)
        models[0].params = p_one
        res_one = backtest(models, code, st, et, file_name=None, res_type='res')
        res_all.append(res_one)
    par_all = pd.DataFrame(par_all)
    res_all = pd.DataFrame(res_all)
    res_all = pd.concat([par_all, res_all], axis=1)
    res_all.to_excel(f'./all_param/all_param_{code}_{st}_{et}_ga.xlsx')


def select_params2(models: list[DNGNode, ...], code: str, st: str, et: str) -> None:
    """
    搜参系统， 采用遍历
    :param models: 以DNGNode列表组成的策略
    :param code: 回测标的代码
    :param st: 开始时间
    :param et: 结束时间
    :param n:运行轮次
    :return: None
    """
    #  初始种群
    p = copy.copy(models[0].__params__)
    res = copy.copy(models[0].__params__)
    for k, v in p.items():
        p[k] = analysis_param(v)
    p_all = ', '.join(map(str, p.values()))
    p_all = eval(f'product({p_all})')
    for k, v in res.items():
        res[k] = []

    fp = ['net', 'avg', 'std', 'shape', 'maxDD', 'calmar']
    for r in fp:
        res[r] = []

    for i in p_all:
        c = 0
        for pp in p:
            p[pp] = i[c]
            res[pp].append(i[c])
            c += 1
        models[0].params = p
        rrr = backtest(models, code, st, et, file_name=None, res_type='res')
        for rr in rrr:
            res[rr].append(rrr[rr])
    res = pd.DataFrame(res)
    res.to_excel(f'./all_param/all_param_{code}_{st}_{et}_all.xlsx')


def rolling_back(models: list[DNGNode, ...], code: str, st: str, et: str,
                 file_name: str = 'roll', gap: int = 365) -> None:
    """
    滚动回测系统1， 返回每一年的状况
    :param models: 以DNGNode列表组成的策略
    :param code:标的代码
    :param st: 开始时间
    :param et: 结束时间
    :param file_name: 储存文件夹名
    :param gap: 回测间隔天数，默认为365天
    :return: None
    """
    date1 = datetime.datetime.strptime(st, '%Y%m%d')
    date2 = datetime.datetime.strptime(et, '%Y%m%d')
    de1 = date1
    de2 = date1 + datetime.timedelta(days=gap)
    if not os.path.exists(file_name):
        os.mkdir(file_name)
    else:
        rmtree(file_name)
        os.mkdir(file_name)
    while de2 < date2:
        print(de1, de2)
        backtest(models, code, datetime.datetime.strftime(de1, '%Y%m%d'),
                 datetime.datetime.strftime(de2, '%Y%m%d'),
                 file_name=f'{file_name}/{datetime.datetime.strftime(de1, "%Y%m%d")}_' + \
                           f'{datetime.datetime.strftime(de2, "%Y%m%d")}.xlsx')
        de1 = de1 + datetime.timedelta(days=gap)

        de2 = de1 + datetime.timedelta(days=gap)
        if de2 > date2:
            de2 = date2

    path = os.listdir(file_name)
    pa = pd.DataFrame()
    for p in path:
        p_ = pd.read_excel(f'{file_name}/' + p)
        pa = pd.concat([pa, p_], axis=0)
    pa = pa.sort_values('Unnamed: 0').set_index('Unnamed: 0')
    pa.to_excel(f'{file_name}/' + 'res.xlsx')


def rolling_back2(models: list[DNGNode, ...], code: str, st: str = None, et: str = None,
                  day_num: int = 252, file_name: str = None) -> None:
    """
    滚动回测2， 在每一天开始回测一段时间
    :param models: 以DNGNode列表组成的策略
    :param code: 回测标的代码
    :param st: 开始时间
    :param et: 结束时间
    :param day_num: 回测间隔，为交易天数
    :return: None
    """
    de = ChoicesData(code, st, et).date.tolist()
    de.sort()
    net_all = []
    date_all = []
    avg_all = []
    shape_all = []
    std_all = []
    maxdd_all = []
    calmar_all = []
    for i in range(len(de) - day_num):
        r = backtest(models, code, datetime.datetime.strftime(de[i], '%Y%m%d'),
                     datetime.datetime.strftime(de[i + day_num - 1], '%Y%m%d'),
                     res_type='res')
        date_all.append(de[i])
        net_all.append(r['net'])
        avg_all.append(r['avg'])
        shape_all.append(r['shape'])
        std_all.append(r['std'])
        maxdd_all.append(r['maxDD'])
        calmar_all.append(r['calmar'])
    p = pd.DataFrame({'date': date_all, 'net': net_all, 'avg': avg_all, 'std': std_all, 'shape': shape_all,
                      'maxDD': maxdd_all, 'calmar': calmar_all})
    if file_name is None:
        file_name = f'roll_day_rr_{code}_{st}_{et}.xlsx'
    p.to_excel(file_name)

def summarization_data(models: list[DNGNode, ...], code: str, st: str, et: str, sel_type: str = 'all'):
    if not os.path.exists(code):
        os.mkdir(code)
    else:
        rmtree(code)
        os.mkdir(code)
    writer = pd.ExcelWriter(f'{code}/{code}_{st}_{et}.xlsx')
    if sel_type == 'all':
        if not os.path.exists(f'./all_param/all_param_{code}_{st}_{et}_all.xlsx'):
            select_params2(models, code, st, et)
            res_par = pd.read_excel(f'./all_param/all_param_{code}_{st}_{et}_all.xlsx')
        else:
            res_par = pd.read_excel(f'./all_param/all_param_{code}_{st}_{et}_all.xlsx')
    elif sel_type == 'ga':
        if not os.path.exists(f'./all_param/all_param_{code}_{st}_{et}_ga.xlsx'):
            select_params(models, code, st, et)
            res_par = pd.read_excel(f'./all_param/all_param_{code}_{st}_{et}_ga.xlsx')
        else:
            res_par = pd.read_excel(f'./all_param/all_param_{code}_{st}_{et}_ga.xlsx')
    else:
        raise
    res_par['sco'] = res_par['shape']
    res_par = res_par.sort_values('sco', ascending=False)
    net = pd.DataFrame()
    counts = []
    res_all = []
    for i in range(0, 5):
        # i = 0
        sel = res_par.iloc[i]
        p = models[0].params
        for j in p:
            p[j] = sel[j]
        models[0].params = p
        res_one = backtest(models, code, st, et, plot=False, file_name='linshi.xlsx', res_type='res')
        res_all.append(res_one)
        res_linshi = pd.read_excel('linshi.xlsx')
        if len(net) == 0:
            net = res_linshi.loc[:, ['Unnamed: 0', 'line_rnt']]
            net.columns = ['Unnamed: 0', f'{i}_{str(p)}']
            net = net.set_index('Unnamed: 0')
        else:
            li = res_linshi.loc[:, ['Unnamed: 0', 'line_rnt']]
            li.columns = ['Unnamed: 0', f'{i}_{str(p)}']
            net = pd.merge(net, li, left_index=True, right_on='Unnamed: 0')
            net = net.set_index('Unnamed: 0')
        count = np.where(res_linshi.loc[:, 'charge'] != 0, 1, 0).sum() / len(res_linshi) * 252
        counts.append(count)
    res_ind = backtest(model_ind, code, st, et, plot=False, file_name='linshi.xlsx', res_type='res')
    res_all.append(res_ind)
    counts.append(1)
    res_linshi = pd.read_excel('linshi.xlsx')
    li = res_linshi.loc[:, ['Unnamed: 0', 'line_rnt']]
    li.columns = ['Unnamed: 0', f'index_']
    net = pd.merge(net, li, left_index=True, right_on='Unnamed: 0')
    net = net.set_index('Unnamed: 0')
    res_all = pd.DataFrame(res_all)
    res_all.loc[:, 'count_num'] = counts
    net.to_excel(writer, sheet_name='net')
    res_all.to_excel(writer, sheet_name='res')
    sel = res_par.iloc[0]
    p = models[0].params
    for j in p:
        p[j] = sel[j]
    models[0].params = p
    roll_res = net.loc[:, []]
    for n in [252, int(252 / 2)]:
        rolling_back2(models, code, st, et, day_num=n, file_name='linshi.xlsx')
        roll_linshi = pd.read_excel('linshi.xlsx')
        li = roll_linshi.loc[:, ['date', 'net']]
        li.columns = ['date', f'strategy_{n}']
        roll_res = pd.merge(roll_res, li, left_index=True, right_on='date', how='left')
        roll_res = roll_res.set_index('date')
        rolling_back2(model_ind, code, st, et, day_num=n, file_name='linshi.xlsx')
        roll_linshi_index = pd.read_excel('linshi.xlsx')
        li = roll_linshi_index.loc[:, ['date', 'net']]
        li.columns = ['date', f'index_{n}']
        roll_res = pd.merge(roll_res, li, left_index=True, right_on='date', how='left')
        roll_res = roll_res.set_index('date')
    roll_res.to_excel(writer, sheet_name='roll_res')
    writer.save()
    writer.close()



